import datetime
import backtrader as bt
import akshare as ak
import pandas as pd


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


def run_backtrade():
    stocks = ['sz002596', 'sz000001']
    cerebro = bt.Cerebro()
    cerebro.addstrategy(TestStrategy)
    st_date = datetime.datetime(2023, 1, 1)
    ed_date = datetime.datetime(2023, 8, 6)

    for i, code in enumerate(stocks):
        print(f"i={i}->{code} Done !")
        data = get_data(code)
        datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_date, todate=ed_date)
        cerebro.adddata(datafeed, name=code)
        print(f'{code} feeds ok')
        break

    # 初始资金 1,000,000.0
    cerebro.broker.setcash(1e6)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    strat = result[0]
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.plot()


class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
        ('printlog', True),
    )

    def __init__(self):
        self.volume = self.datas[0].volume
        self.daily_pnl = dict()
        self.trade_info = list()
        self.signal = dict()
        self.buy_info = dict()
        self.order = dict()
        self.buyprice = None
        self.buycomm = None

        for data in self.datas:
            dif = bt.indicators.EMA(data, period=12, subplot=False) - bt.indicators.EMA(data, period=26, subplot=False)
            dea = bt.indicators.EMA(dif, period=9)
            # self.signal[data] = dif - dea
            self.signal[data] = bt.ind.CrossOver(dif, dea)

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print('%s, %s' % (dt.isoformat(), txt))

    @staticmethod
    def percent(today, yesterday):
        return float(today - yesterday) / today

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    f'code={order.data._name} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={order.data._name} BUY EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f' Comm {order.executed.value:.2f}, Size: {order.executed.size:.2f}')

                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
                self.buy_info[order.data._name] = order.executed.price
            else:
                self.log(
                    f'code={order.data._name} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={order.data._name} SELL EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f' Comm {order.executed.value:.2f}, Size: {order.executed.size:.2f}')
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        self.order[order.data._name] = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {trade.pnl:.2f}, NET {trade.pnlcomm:.2f}')

    def next(self):
        for data in self.datas:
            code = data._name
            pos = self.getposition(data).size
            print(f'--->code={code}->{data.datetime.datetime(0)}->pos={pos}')

            # if self.signal[data](-1) < 0 and self.signal[data](0) > 0:
            #    self.order = self.buy()
            if not pos:
                if self.signal[data] > 0:
                    print(
                        f'--->code={code}->{data.datetime.datetime(0)}->signal={self.signal[data][0]}')
                    self.order[code] = self.buy(data=data)

            elif code in self.buy_info:
                condition = data.close[0] / self.buy_info[code] - 1
                print(f'--->code={code}->{data.datetime.datetime(0)}->close={data.close[0]}->cost={self.buy_info[code]}->signal={self.signal[data][0]}->condition={condition}')
                if (self.signal[data] < 0) and (condition > 0.1 or condition < -0.1):
                    self.order[code] = self.sell(data=data)
                    # del self.buy_info[code]


if __name__ == '__main__':
    run_backtrade()
